/*
 * Copyright (C) 2010-2101 Alibaba Group Holding Limited.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.alibaba.otter.node.etl.load.loader.db;

import com.alibaba.otter.node.etl.common.db.dialect.DbDialect;
import com.alibaba.otter.node.etl.common.db.dialect.DbDialectFactory;
import com.alibaba.otter.node.etl.common.db.dialect.mysql.MysqlDialect;
import com.alibaba.otter.node.etl.common.db.utils.SqlUtils;
import com.alibaba.otter.node.etl.load.exception.LoadException;
import com.alibaba.otter.node.etl.load.loader.AbstractLoadAction;
import com.alibaba.otter.node.etl.load.loader.db.context.DataLoadContext;
import com.alibaba.otter.shared.common.model.config.ConfigHelper;
import com.alibaba.otter.shared.common.model.config.data.DataMedia;
import com.alibaba.otter.shared.common.model.config.data.db.DbMediaSource;
import com.alibaba.otter.shared.etl.model.EventColumn;
import com.alibaba.otter.shared.etl.model.EventData;
import com.alibaba.otter.shared.etl.model.EventType;
import com.alibaba.otter.shared.etl.model.ObjectData;
import org.apache.commons.lang.StringUtils;
import org.apache.ddlutils.model.Column;
import org.apache.ddlutils.model.Table;
import org.springframework.jdbc.core.BatchPreparedStatementSetter;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.jdbc.core.StatementCallback;
import org.springframework.jdbc.core.StatementCreatorUtils;
import org.springframework.jdbc.support.lob.LobCreator;
import org.springframework.util.CollectionUtils;

import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.sql.Types;
import java.util.*;
import java.util.concurrent.Callable;
import java.util.stream.Collectors;

/**
 * 数据库load的执行入口
 *
 * @author jianghang 2011-10-31 下午03:17:43
 * @version 4.0.0
 */
public class DbLoadAction extends AbstractLoadAction {


    private DbDialectFactory dbDialectFactory;

    public DbLoadAction() {
        this.workName = "DbLoadAction";
    }

    /**
     * 执行ddl的调用，处理逻辑比较简单: 串行调用
     */
    @Override
    protected boolean doLoadDDL(DataLoadContext context, EventData eventData) {

        DataMedia<?> dataMedia = ConfigHelper.findDataMedia(context.getPipeline(), eventData.getTableId());
        final DbDialect dbDialect = dbDialectFactory.getDbDialect(context.getIdentity().getPipelineId(),
                (DbMediaSource) dataMedia.getSource());
        // run ddl
        // 如果mysql，执行ddl时，切换到在源库执行的schema上
        // result &= stmt.execute("use " +
        // data.getDdlSchemaName());
        // 解决当数据库名称为关键字如"Order"的时候,会报错,无法同步
        return Boolean.TRUE.equals(dbDialect.getJdbcTemplate().execute((StatementCallback<Boolean>) stmt -> {
            boolean result1 = true;
            if (dbDialect instanceof MysqlDialect && StringUtils.isNotEmpty(eventData.getDdlSchemaName())) {
                // 如果mysql，执行ddl时，切换到在源库执行的schema上
                // result &= stmt.execute("use " +
                // data.getDdlSchemaName());

                // 解决当数据库名称为关键字如"Order"的时候,会报错,无法同步
                result1 = stmt.execute("use `" + eventData.getDdlSchemaName() + "`");
            }
            result1 &= stmt.execute(eventData.getSql());
            return result1;
        }));

    }

    @Override
    protected void load(DataLoadContext context, List<EventData> items) {
        //判断是否启用数据行合并功能
        // 进行一次数据合并，合并相同pk的多次I/U/D操作
        items = DbLoadMerger.merge(items);
        // 按I/U/D进行归并处理
        DbLoadData loadData = new DbLoadData();
        doBefore(items, context, loadData);
        // 执行load操作
        doLoad(context, loadData);
    }

    private void doBefore(List<EventData> items, DataLoadContext context, DbLoadData loadData) {
        for (final EventData item : items) {
            boolean filter = interceptor.before(context, item);
            if (!filter) {
                loadData.merge(item);// 进行分类
            }
        }
    }



    private void doLoad(final DataLoadContext context, DbLoadData loadData) {
        // 优先处理delete,可以利用batch优化
        List<List<EventData>> batchDatas = new ArrayList<>();
        List<DbLoadData.TableLoadData> tables = loadData.getTables();
        // 优先处理delete,可以利用batch优化
        for (DbLoadData.TableLoadData tableData : tables) {
            if (useBatch) {
                // 优先执行delete语句，针对uniqe更新，一般会进行delete + insert的处理模式，避免并发更新
                batchDatas.addAll(split(tableData.getDeleteDatas()));
            } else {
                // 如果不可以执行batch，则按照单条数据进行并行提交
                // 优先执行delete语句，针对uniqe更新，一般会进行delete + insert的处理模式，避免并发更新
                for (EventData data : tableData.getDeleteDatas()) {
                    batchDatas.add(Collections.singletonList(data));
                }
            }
        }

        if (context.getPipeline().getParameters().isDryRun()) {
            doDryRun(context, batchDatas);
        } else {
            doTwoPhase(context, batchDatas);
        }
        batchDatas.clear();

        // 处理下insert/update
        for (DbLoadData.TableLoadData tableData : tables) {
            if (useBatch) {
                // 执行insert + update语句
                batchDatas.addAll(split(tableData.getInsertDatas()));
                batchDatas.addAll(split(tableData.getUpadateDatas()));// 每条记录分为一组，并行加载
            } else {
                // 执行insert + update语句
                for (EventData data : tableData.getInsertDatas()) {
                    batchDatas.add(Collections.singletonList(data));
                }
                for (EventData data : tableData.getUpadateDatas()) {
                    batchDatas.add(Collections.singletonList(data));
                }
            }
        }

        if (context.getPipeline().getParameters().isDryRun()) {
            doDryRun(context, batchDatas);
        } else {
            doTwoPhase(context, batchDatas);
        }
        batchDatas.clear();
    }

    /**
     * 将对应的数据按照sql相同进行batch组合
     */
    private List<List<EventData>> split(List<EventData> datas) {

        List<List<EventData>> result = new ArrayList<>();
        if (CollectionUtils.isEmpty(datas)) {
            return result;
        }

        int[] bits = new int[datas.size()];// 初始化一个标记，用于标明对应的记录是否已分入某个batch
        for (int i = 0; i < bits.length; i++) {
            // 跳过已经被分入batch的
            while (i < bits.length && bits[i] == 1) {
                i++;
            }

            if (i >= bits.length) { // 已处理完成，退出
                break;
            }

            // 开始添加batch，最大只加入batchSize个数的对象
            List<EventData> batch = new ArrayList<>();
            bits[i] = 1;
            batch.add(datas.get(i));
            for (int j = i + 1; j < bits.length && batch.size() < batchSize; j++) {
                if (bits[j] == 0 && canBatch(datas.get(i), datas.get(j))) {
                    batch.add(datas.get(j));
                    bits[j] = 1;// 修改为已加入
                }
            }
            result.add(batch);
        }

        return result;

    }

    /**
     * 判断两条记录是否可以作为一个batch提交，主要判断sql是否相等. 可优先通过schemaName进行判断
     */
    private boolean canBatch(EventData source, EventData target) {

        // return StringUtils.equals(source.getSchemaName(),
        // target.getSchemaName())
        // && StringUtils.equals(source.getTableName(), target.getTableName())
        // && StringUtils.equals(source.getSql(), target.getSql());
        // return StringUtils.equals(source.getSql(), target.getSql());

        // 因为sqlTemplate构造sql时用了String.intern()的操作，保证相同字符串的引用是同一个，所以可以直接使用==进行判断，提升效率
        return source.getSql() == target.getSql();
    }


    @Override
    protected Callable<Exception> getWorker(DataLoadContext context, List<EventData> eventDataCollection, boolean canBatch) {
        return new DbLoadWorker(context, eventDataCollection, canBatch);
    }

    class DbLoadWorker extends AbstractWorker{

        private final static String WORK_NAME = "pipelineId = %s,pipelineName = %s ,%s";

        private final DbDialect dbDialect;

        public DbLoadWorker(DataLoadContext context, List<EventData> eventDataList, boolean canBatch) {
            super(context, eventDataList, canBatch);
            EventData data = eventDataList.get(0); // eventData为同一数据库的记录，只取第一条即可
            DataMedia<?> dataMedia = ConfigHelper.findDataMedia(context.getPipeline(), data.getTableId());
            dbDialect = dbDialectFactory.getDbDialect(context.getIdentity().getPipelineId(),
                    (DbMediaSource) dataMedia.getSource());
        }

        @Override
        protected String getThreadName() {
            return String.format(WORK_NAME,context.getPipeline().getId(),context.getPipeline().getName(),"db-worker");
        }

        @Override
        protected void doLoad(EventData data) {
            final LobCreator lobCreator = dbDialect.getLobHandler().getLobCreator();
            int affect = dbDialect.getTransactionTemplate().execute(status -> {
                try {
                    failedDataList.clear(); // 先清理
                    processedDataList.clear();
                    interceptor.transactionBegin(context, Collections.singletonList(data), dbDialect);
                    JdbcTemplate template = dbDialect.getJdbcTemplate();
                    int affect1 = template.update(data.getSql(), ps -> doPreparedStatement(ps, dbDialect, lobCreator, data));
                    interceptor.transactionEnd(context, Collections.singletonList(data), dbDialect);
                    return affect1;
                } finally {
                    lobCreator.close();
                }
            });
            // 更新统计信息
            processStat(data, affect, false);
        }

        @Override
        protected void doBatchLoad(List<EventData> splitDataList) {

            final LobCreator lobCreator = dbDialect.getLobHandler().getLobCreator();
            final String sql = splitDataList.get(0).getSql();
            int[] affects = dbDialect.getTransactionTemplate().execute(status -> {
                // 初始化一下内容
                try {
                    failedDataList.clear(); // 先清理
                    processedDataList.clear();
                    List<ObjectData> collect = splitDataList.stream().map(ObjectData.class::cast).collect(Collectors.toList());
                    interceptor.transactionBegin(context, collect, dbDialect);
                    JdbcTemplate template = dbDialect.getJdbcTemplate();
                    int[] affects1 = template.batchUpdate(sql, new BatchPreparedStatementSetter() {

                        public void setValues(PreparedStatement ps, int idx) throws SQLException {
                            doPreparedStatement(ps, dbDialect, lobCreator, splitDataList.get(idx));
                        }

                        public int getBatchSize() {
                            return splitDataList.size();
                        }
                    });
                    interceptor.transactionEnd(context, collect, dbDialect);
                    return affects1;
                } finally {
                    lobCreator.close();
                }
            });

            // 更新统计信息
            for (int i = 0; i < splitDataList.size(); i++) {
                processStat(splitDataList.get(i), affects[i], true);
            }
        }

        private void doPreparedStatement(PreparedStatement ps, DbDialect dbDialect, LobCreator lobCreator,
                                         EventData data) throws SQLException {
            EventType type = data.getEventType();
            // 注意insert/update语句对应的字段数序都是将主键排在后面
            List<EventColumn> columns = new ArrayList<>();
            if (type.isInsert()) {
                columns.addAll(data.getColumns()); // insert为所有字段
                columns.addAll(data.getKeys());
            } else if (type.isDelete()) {
                columns.addAll(data.getKeys());
            } else if (type.isUpdate()) {
                boolean existOldKeys = !CollectionUtils.isEmpty(data.getOldKeys());
                columns.addAll(data.getUpdatedColumns());// 只更新带有isUpdate=true的字段
                if (existOldKeys && dbDialect.isDRDS()) {
                    // DRDS需要区分主键是否有变更
                    columns.addAll(data.getUpdatedKeys());
                } else {
                    columns.addAll(data.getKeys());
                }
                if (existOldKeys) {
                    columns.addAll(data.getOldKeys());
                }
            }

            // 获取一下当前字段名的数据是否必填
            Table table = dbDialect.findTable(data.getSchemaName(), data.getTableName());
            Map<String, Boolean> isRequiredMap = new HashMap<>();
            for (Column tableColumn : table.getColumns()) {
                isRequiredMap.put(StringUtils.lowerCase(tableColumn.getName()), tableColumn.isRequired());
            }

            for (int i = 0; i < columns.size(); i++) {
                int paramIndex = i + 1;
                EventColumn column = columns.get(i);
                int sqlType = column.getColumnType();

                Boolean isRequired = isRequiredMap.get(StringUtils.lowerCase(column.getColumnName()));
                if (isRequired == null) {
                    // 清理一下目标库的表结构,二次检查一下
                    table = dbDialect.findTable(data.getSchemaName(), data.getTableName(), false);
                    isRequiredMap = new HashMap<>();
                    for (Column tableColumn : table.getColumns()) {
                        isRequiredMap.put(StringUtils.lowerCase(tableColumn.getName()), tableColumn.isRequired());
                    }

                    isRequired = isRequiredMap.get(StringUtils.lowerCase(column.getColumnName()));
                    if (isRequired == null) {
                        throw new LoadException(String.format("column name %s is not found in Table[%s]",
                                column.getColumnName(), table));
                    }
                }

                Object param;
                if (dbDialect instanceof MysqlDialect
                        && (sqlType == Types.TIME || sqlType == Types.TIMESTAMP || sqlType == Types.DATE)) {
                    // 解决mysql的0000-00-00 00:00:00问题，直接依赖mysql
                    // driver进行处理，如果转化为Timestamp会出错
                    param = column.getColumnValue();
                } else {
                    param = SqlUtils.stringToSqlValue(column.getColumnValue(),
                            sqlType,
                            isRequired,
                            dbDialect.isEmptyStringNulled());
                }

                try {
                    switch (sqlType) {
                        case Types.CLOB:
                            lobCreator.setClobAsString(ps, paramIndex, (String) param);
                            break;

                        case Types.BLOB:
                            lobCreator.setBlobAsBytes(ps, paramIndex, (byte[]) param);
                            break;
                        case Types.TIME:
                        case Types.TIMESTAMP:
                        case Types.DATE:
                            // 只处理mysql的时间类型，oracle的进行转化处理
                            if (dbDialect instanceof MysqlDialect) {
                                // 解决mysql的0000-00-00 00:00:00问题，直接依赖mysql
                                // driver进行处理，如果转化为Timestamp会出错
                                ps.setObject(paramIndex, param);
                            } else {
                                StatementCreatorUtils.setParameterValue(ps, paramIndex, sqlType, null, param);
                            }
                            break;
                        case Types.BIT:
                            // 只处理mysql的bit类型，bit最多存储64位，所以需要使用BigInteger进行处理才能不丢精度
                            // mysql driver将bit按照setInt进行处理，会导致数据越界
                            if (dbDialect instanceof MysqlDialect) {
                                StatementCreatorUtils.setParameterValue(ps, paramIndex, Types.DECIMAL, null, param);
                            } else {
                                StatementCreatorUtils.setParameterValue(ps, paramIndex, sqlType, null, param);
                            }
                            break;
                        default:
                            StatementCreatorUtils.setParameterValue(ps, paramIndex, sqlType, null, param);
                            break;
                    }
                } catch (SQLException ex) {
                    logger.error("## SetParam error , [pairId={}, sqltype={}, value={}]",
                            data.getPairId(), sqlType, param);
                    throw ex;
                }
            }
        }
    }


    // =============== setter / getter ===============
    public void setRetry(int retry) {
        this.retry = retry;
    }

    public void setRetryWait(int retryWait) {
        this.retryWait = retryWait;
    }

    public void setDbDialectFactory(DbDialectFactory dbDialectFactory) {
        this.dbDialectFactory = dbDialectFactory;
    }
}
